package network

import (
	"context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials/insecure"
	"log"
	"sync"
	"time"
)

var (
	senderID     int64
	verbose      bool
	senderWg     sync.WaitGroup
	broadcastTimer int
	sleepTimerValue int
	dialOpt      []grpc.DialOption
	connections  AddrConnMap
)

// AddrConnMap manages gRPC connections
type AddrConnMap struct {
	m      map[string]SendClient
	idmap  map[string]string
	count  map[string]int
	mu     sync.RWMutex
}

func (s *AddrConnMap) Init() {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.m = make(map[string]SendClient)
	s.count = make(map[string]int)
	s.idmap = make(map[string]string)
}

func (s *AddrConnMap) Get(key string) (SendClient, bool) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	conn, exist := s.m[key]
	return conn, exist && conn != nil
}

func (s *AddrConnMap) GetID(key string) string {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.idmap[key]
}

func (s *AddrConnMap) Insert(key string, value SendClient) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.m[key] = value
}

func (s *AddrConnMap) InsertID(key string, value string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.idmap[key] = value
}

// SendClient is the interface for gRPC send client
type SendClient interface {
	RBCSendByteMsg(ctx context.Context, in *RawMessage, opts ...grpc.CallOption) (*Empty, error)
	ABASendByteMsg(ctx context.Context, in *RawMessage, opts ...grpc.CallOption) (*Empty, error)
	PRFSendByteMsg(ctx context.Context, in *RawMessage, opts ...grpc.CallOption) (*Empty, error)
	ECRBCSendByteMsg(ctx context.Context, in *RawMessage, opts ...grpc.CallOption) (*Empty, error)
	CBCSendByteMsg(ctx context.Context, in *RawMessage, opts ...grpc.CallOption) (*Empty, error)
	EVCBCSendByteMsg(ctx context.Context, in *RawMessage, opts ...grpc.CallOption) (*Empty, error)
	MVBASendByteMsg(ctx context.Context, in *RawMessage, opts ...grpc.CallOption) (*Empty, error)
	RetrieveSendByteMsg(ctx context.Context, in *RawMessage, opts ...grpc.CallOption) (*Empty, error)
	CBRSendByteMsg(ctx context.Context, in *RawMessage, opts ...grpc.CallOption) (*Empty, error)
}

// RawMessage and Empty are defined in the proto file
type RawMessage struct {
	Msg    []byte
	Result bool
}

type Empty struct{}

type MessageType int32

const (
	MessageType_BROADCAST   MessageType = 0
	MessageType_JOIN        MessageType = 1
	MessageType_WRITE       MessageType = 2
	MessageType_WRITE_BATCH MessageType = 3
)

type Request struct {
	Type    MessageType
	Request []byte
	Version string
}

func BuildConnection(ctx context.Context, nid string, address string, clientFactory func(*grpc.ClientConn) SendClient) bool {
	if verbose {
		log.Printf("building a connection with %v", nid)
	}

	conn, err := grpc.DialContext(ctx, address, dialOpt...)
	if err != nil {
		log.Printf("[Communication Sender Error] failed to build a connection with %v", err)
		return false
	}

	c := clientFactory(conn)
	connections.Insert(address, c)
	connections.InsertID(address, nid)
	return true
}

// BuildConnectionFromFactory builds connection with proper client factory
func BuildConnectionFromFactory(ctx context.Context, nid string, address string) bool {
	return BuildConnection(ctx, nid, address, CreateGRPCClient)
}

func ByteSend(msg []byte, address string, msgType TypeOfMessage, clientFactory func(*grpc.ClientConn) SendClient) {
	if address == "" {
		return
	}

	ctx, cancel := context.WithTimeout(context.Background(), time.Duration(broadcastTimer)*time.Millisecond)
	defer cancel()

	// Get node ID from address (simplified - should map address to node ID properly)
	nid := Int64ToString(senderID) // This should be determined from address mapping
	c, built := connections.Get(address)
	existnid := connections.GetID(address)

	if !built || c == nil || nid != existnid {
		// Build connection
		suc := BuildConnectionFromFactory(ctx, nid, address)
		if !suc {
			log.Printf("[Communication Sender Error] did not connect to node at %s, set it to notlive", address)
			NotLive(nid)
			broadcastTimer = broadcastTimer * 2
			return
		}
		c, _ = connections.Get(address)
	}

	if c == nil {
		log.Printf("[Communication Sender Error] failed to get connection for %s", address)
		return
	}

	var err error
	switch msgType {
	case RBC_ALL:
		_, err = c.RBCSendByteMsg(ctx, &RawMessage{Msg: msg})
	case ABA_ALL:
		_, err = c.ABASendByteMsg(ctx, &RawMessage{Msg: msg})
	case PRF:
		_, err = c.PRFSendByteMsg(ctx, &RawMessage{Msg: msg})
	case ECRBC_ALL:
		_, err = c.ECRBCSendByteMsg(ctx, &RawMessage{Msg: msg})
	case CBC_ALL:
		_, err = c.CBCSendByteMsg(ctx, &RawMessage{Msg: msg})
	case EVCBC_ALL:
		_, err = c.EVCBCSendByteMsg(ctx, &RawMessage{Msg: msg})
	case MVBA_DISTRIBUTE:
		_, err = c.MVBASendByteMsg(ctx, &RawMessage{Msg: msg})
	case RETRIEVE:
		_, err = c.RetrieveSendByteMsg(ctx, &RawMessage{Msg: msg})
	case CBR_ALL:
		_, err = c.CBRSendByteMsg(ctx, &RawMessage{Msg: msg})
	default:
		log.Printf("message type %v not supported", msgType)
		return
	}

	if err != nil {
		log.Printf("[Communication Sender Error] could not send message to node %s: %v", nid, err)
		NotLive(nid)
		connections.Insert(address, nil)
		return
	}
}

func MACBroadcast(msg []byte, mtype ProtocolType, addresses []string, clientFactory func(*grpc.ClientConn) SendClient) {
	for _, address := range addresses {
		// Get node ID from address (simplified)
		nid := Int64ToString(senderID) // Should map address to node ID
		dest, _ := StringToInt64(nid)

		request, err := SerializeWithMAC(senderID, dest, msg)
		if err != nil {
			log.Printf("[Sender Error] Not able to generate MAC")
			continue
		}

		if IsNotLive(nid) {
			if verbose {
				log.Printf("[Communication Sender] Replica %v is not live, don't send message to it", nid)
			}
			continue
		}

		var msgType TypeOfMessage
		switch mtype {
		case RBC:
			msgType = RBC_ALL
		case ABA:
			msgType = ABA_ALL
		case CBC:
			msgType = CBC_ALL
		case EVCBC:
			msgType = EVCBC_ALL
		case MVBA:
			msgType = MVBA_DISTRIBUTE
		case CBR:
			msgType = CBR_ALL
		default:
			msgType = CBR_ALL
		}

		go ByteSend(request, address, msgType, clientFactory)
	}
}

func SendToNode(msg []byte, dest int64, mtype ProtocolType, address string, clientFactory func(*grpc.ClientConn) SendClient) {
	nid := Int64ToString(dest)

	request, err := SerializeWithMAC(senderID, dest, msg)
	if err != nil {
		log.Printf("[Sender Error] Not able to generate MAC")
		return
	}

	if IsNotLive(nid) {
		if verbose {
			log.Printf("[Communication Sender] Replica %v is not live, don't send message to it", nid)
		}
		return
	}

	var msgType TypeOfMessage
	switch mtype {
	case CBC:
		msgType = CBC_ALL
	case EVCBC:
		msgType = EVCBC_ALL
	case CBR:
		msgType = CBR_ALL
	default:
		msgType = CBR_ALL
	}

	go ByteSend(request, address, msgType, clientFactory)
}

func StartSender(rid int64, addrs []string) {
	log.Printf("Starting sender %v", rid)
	senderID = rid
	verbose = true

	dialOpt = []grpc.DialOption{
		grpc.WithTransportCredentials(insecure.NewCredentials()),
		grpc.WithBlock(),
	}

	connections.Init()
	StartConnectionManager()
	broadcastTimer = 5000 // 5 seconds default
	sleepTimerValue = 1000
}

func SetId(newnid int64) {
	senderID = newnid
}

